from lazyllm import LOG
import re
from typing import Optional, Dict
MODEL_MAPPING = {
    # ===== OpenAI (LLM) =====
    'gpt-5': 'vlm',
    'gpt-5-mini': 'vlm',
    'gpt-5-nano': 'vlm',
    'gpt-5-codex': 'vlm',
    'gpt-5-pro': 'llm',
    'sora-2': 'sd',
    'sora-2-pro': 'sd',
    'gpt-image-1': 'sd',
    'gpt-image-1-mini': 'sd',
    'o3-deep-research': 'llm',
    'o4-mini-deep-research': 'llm',
    'o3-pro': 'llm',
    'gpt-audio': 'llm',
    'gpt-realtime': 'llm',
    'gpt-realtime-mini': 'llm',
    'gpt-audio-mini': 'llm',
    'o3': 'vlm',
    'o4-mini': 'vlm',
    'gpt-4.1': 'vlm',
    'gpt-4.1-mini': 'vlm',
    'gpt-4.1-nano': 'vlm',
    'o1-pro': 'vlm',
    'computer-use-preview': 'vlm',
    'gpt-4o-mini-search-preview': 'llm',
    'gpt-4o-search-preview': 'llm',
    'o3-mini': 'llm',
    'gpt-4o-mini-audio-preview': 'llm',
    'gpt-4o-mini-realtime-preview': 'llm',
    'o1': 'llm',
    'omni-moderation-latest': 'vlm',
    'gpt-4o': 'vlm',
    'gpt-4o-audio-preview': 'llm',
    'gpt-4o-mini': 'vlm',
    'gpt-4o-realtime-preview': 'llm',
    'gpt-4-turbo': 'vlm',
    'chatgpt-4o-latest': 'vlm',
    'codex-mini-latest': 'vlm',
    'dall-e-2': 'sd',
    'dall-e-3': 'sd',
    'gpt-3.5-turbo': 'llm',
    'gpt-4': 'llm',
    'gpt-4o-mini-transcribe': 'stt',
    'gpt-4o-mini-tts': 'tts',
    'gpt-4o-transcribe': 'stt',
    'gpt-4o-transcribe-diarize': 'stt',
    'gpt-5-chat-latest': 'vlm',
    'gpt-oss-120b': 'llm',
    'gpt-oss-20b': 'llm',
    'text-embedding-3-large': 'embed',
    'text-embedding-3-small': 'embed',
    'text-embedding-ada-002': 'embed',
    'tts-1': 'tts',
    'tts-1-hd': 'tts',
    'whisper-1': 'stt',

    # ===== SenseNova =====
    'sensenova-v6-5-pro': 'vlm',
    'sensenova-v6-5-turbo': 'vlm',
    'sensenova-v6-pro': 'vlm',
    'sensenova-v6-turbo': 'vlm',
    'sensenova-v6-reasoner': 'vlm',
    'sensenova-v6-omni': 'vlm',
    'sensechat-5-1202': 'llm',
    'sensechat-turbo-1202': 'llm',
    'sensechat-5': 'llm',
    'sensechat': 'llm',
    'sensechat-32k': 'llm',
    'sensechat-128k': 'llm',
    'sensechat-turbo': 'llm',
    'sensechat-5-cantonese': 'llm',
    'sensechat-character-pro': 'llm',
    'sensechat-character': 'llm',
    'sensechat-vision': 'vlm',
    'sensenova-audio-fusion-0603': 'tts',
    'nova-tts-1': 'tts',
    'nova-embedding-stable': 'embed',

    # ===== GLM =====
    'chatglm3-6b': 'llm',
    'chatglm_12b': 'llm',
    'chatglm_32b': 'llm',
    'chatglm_66b': 'llm',
    'chatglm_130b': 'llm',
    'glm-4-6': 'llm',
    'glm-4-5': 'llm',
    'glm-4-5-x': 'llm',
    'glm-4-5-air': 'llm',
    'glm-4-5-airx': 'llm',
    'glm-4-plus': 'llm',
    'glm-4-air-250414': 'llm',
    'glm-4-long': 'llm',
    'glm-4-airx': 'llm',
    'glm-4-flashx-250414': 'llm',
    'glm-4-5-flash': 'llm',
    'glm-4-flash-250414': 'llm',
    'glm-4-5v': 'vlm',
    'glm-4-1v-thinking-flashx': 'vlm',
    'glm-4v-plus-0111': 'vlm',
    'glm-4-1v-thinking-flash': 'vlm',
    'glm-4v-flash': 'vlm',
    'cogview-4': 'sd',
    'cogview-3-flash': 'sd',
    'cogvideox-3': 'sd',
    'cogvideox-2': 'sd',
    'vidu-q1': 'sd',
    'vidu-2': 'sd',
    'cogvideox-flash': 'sd',
    'glm-4-voice': 'llm',
    'glm-realtime': 'llm',
    'glm-asr': 'stt',
    'embedding-3': 'embed',
    'embedding-2': 'embed',
    'charglm-4': 'llm',
    'emohaa': 'llm',
    'codegeex-4': 'llm',
    'rerank': 'rerank',

    # ===== Kimi (Moonshot) =====
    'kimi-k2-0905-preview': 'llm',
    'kimi-k2-0711-preview': 'llm',
    'kimi-k2-turbo-preview': 'llm',
    'moonshot-v1-8k': 'llm',
    'moonshot-v1-32k': 'llm',
    'moonshot-v1-128k': 'llm',
    'moonshot-v1-8k-vision-preview': 'vlm',
    'moonshot-v1-32k-vision-preview': 'vlm',
    'moonshot-v1-128k-vision-preview': 'vlm',
    'kimi-latest': 'vlm',
    'kimi-thinking-preview': 'vlm',

    # ===== Qwen =====
    # LLM
    'qwen3-max': 'llm',
    'qwen3-max-2025-09-23': 'llm',
    'qwen3-max-preview': 'llm',
    'qwen-turbo': 'llm',
    'qwen-turbo-latest': 'llm',
    'qwen-turbo-2025-07-15': 'llm',
    'qwen-turbo-2025-04-28': 'llm',
    'qwen-7b-chat': 'llm',
    'qwen-72b-chat': 'llm',
    'qwen-plus': 'llm',
    'qwen-plus-latest': 'llm',
    'qwen-plus-2025-09-11': 'llm',
    'qwen-plus-2025-07-28': 'llm',
    'qwen-plus-2025-07-14': 'llm',
    'qwen-plus-2025-04-28': 'llm',
    'qwen-flash': 'llm',
    'qwen-flash-2025-07-28': 'llm',
    'qwen-long': 'llm',
    'qwen-long-latest': 'llm',
    'qwen-long-2025-01-25': 'llm',
    'qwq-plus': 'llm',
    'qwq-plus-latest': 'llm',
    'qwq-plus-2025-03-05': 'llm',
    'qwen3-omni-flash': 'llm',
    'qwen3-omni-flash-2025-09-15': 'llm',
    'qwen3-omni-flash-realtime': 'llm',
    'qwen3-omni-flash-realtime-2025-09-15': 'llm',
    'qwen-math-plus': 'llm',
    'qwen-math-turbo': 'llm',
    'qwen3-coder-plus': 'llm',
    'qwen3-coder-plus-2025-09-23': 'llm',
    'qwen3-coder-plus-2025-07-22': 'llm',
    'qwen3-coder-flash': 'llm',
    'qwen3-coder-flash-2025-07-28': 'llm',
    'qwen-mt-plus': 'llm',
    'qwen-mt-turbo': 'llm',
    'qwen-doc-turbo': 'llm',
    'qwen-deep-research': 'llm',
    'qwen3-next-80b-a3b-thinking': 'llm',
    'qwen3-next-80b-a3b-instruct': 'llm',
    'qwen3-235b-a22b-thinking-2507': 'llm',
    'qwen3-235b-a22b-instruct-2507': 'llm',
    'qwen3-30b-a3b-thinking-2507': 'llm',
    'qwen3-30b-a3b-instruct-2507': 'llm',
    'qwen3-235b-a22b': 'llm',
    'qwen3-32b': 'llm',
    'qwen3-30b-a3b': 'llm',
    'qwen3-14b': 'llm',
    'qwen3-8b': 'llm',
    'qwen3-4b': 'llm',
    'qwen3-1.7b': 'llm',
    'qwen3-0.6b': 'llm',
    'qwq-32b': 'llm',
    'qwq-32b-preview': 'llm',
    'qwen2.5-14b-instruct-1m': 'llm',
    'qwen2.5-7b-instruct-1m': 'llm',
    'qwen2.5-72b-instruct': 'llm',
    'qwen2.5-32b-instruct': 'llm',
    'qwen2.5-14b-instruct': 'llm',
    'qwen2.5-7b-instruct': 'llm',
    'qwen2.5-3b-instruct': 'llm',
    'qwen2.5-1.5b-instruct': 'llm',
    'qwen2.5-0.5b-instruct': 'llm',
    'qwen2-72b-instruct': 'llm',
    'qwen2-57b-a14b-instruct': 'llm',
    'qwen2-7b-instruct': 'llm',
    'qwen2-1.5b-instruct': 'llm',
    'qwen2-0.5b-instruct': 'llm',
    'qwen1.5-110b-chat': 'llm',
    'qwen1.5-72b-chat': 'llm',
    'qwen1.5-32b-chat': 'llm',
    'qwen1.5-14b-chat': 'llm',
    'qwen1.5-7b-chat': 'llm',
    'qwen1.5-1.8b-chat': 'llm',
    'qwen1.5-0.5b-chat': 'llm',
    'Qwen-Omni': 'llm',
    'qwen2.5-math-72b-instruct': 'llm',
    'qwen2.5-math-7b-instruct': 'llm',
    'qwen2.5-math-1.5b-instruct': 'llm',
    'qwen3-coder-480b-a35b-instruct': 'llm',
    'qwen3-coder-30b-a3b-instruct': 'llm',
    'qwen2.5-coder-32b-instruct': 'llm',
    'qwen2.5-coder-14b-instruct': 'llm',
    'qwen2.5-coder-7b-instruct': 'llm',
    'qwen2.5-coder-3b-instruct': 'llm',
    'qwen2.5-coder-1.5b-instruct': 'llm',
    'qwen2.5-coder-0.5b-instruct': 'llm',

    # VLM
    'qwen-vl-plus': 'vlm',
    'qwen-vl-max': 'vlm',
    'qvq-max': 'vlm',
    'qvq-max-latest': 'vlm',
    'qvq-max-2025-05-15': 'vlm',
    'qvq-plus': 'vlm',
    'qvq-plus-latest': 'vlm',
    'qvq-plus-2025-05-15': 'vlm',
    'qwen3-vl-plus': 'vlm',
    'qwen3-vl-plus-2025-09-23': 'vlm',
    'qwen3-vl-flash': 'vlm',
    'qwen3-vl-flash-2025-10-15': 'vlm',
    'Qwen-VL': 'vlm',
    'qwen3-vl-235b-a22b-thinking': 'vlm',
    'qwen3-vl-235b-a22b-instruct': 'vlm',
    'qwen3-vl-32b-thinking': 'vlm',
    'qwen3-vl-32b-instruct': 'vlm',
    'qwen3-vl-30b-a3b-thinking': 'vlm',
    'qwen3-vl-30b-a3b-instruct': 'vlm',
    'qwen3-vl-8b-thinking': 'vlm',
    'qwen3-vl-8b-instruct': 'vlm',
    'qwen2.5-vl-72b-instruct': 'vlm',
    'qwen2.5-vl-32b-instruct': 'vlm',
    'qwen2.5-vl-7b-instruct': 'vlm',
    'qwen2.5-vl-3b-instruct': 'vlm',
    'qwen2-vl-72b-instruct': 'vlm',
    'qwen2-vl-7b-instruct': 'vlm',
    'qwen2-vl-2b-instruct': 'vlm',
    'qwen-vl-v1': 'vlm',
    'qwen-vl-chat-v1': 'vlm',
    'Qwen3-Omni-Captioner': 'vlm',

    # OCR
    'qwen-vl-ocr': 'ocr',
    'qwen-vl-ocr-latest': 'ocr',
    'qwen-vl-ocr-2025-08-28': 'ocr',
    'qwen-vl-ocr-2025-04-13': 'ocr',
    'qwen-vl-ocr-2024-10-28': 'ocr',

    # STT
    'qwen-audio-turbo': 'stt',
    'qwen-audio-turbo-latest': 'stt',
    'qwen-audio-turbo-2024-12-04': 'stt',
    'qwen3-asr-flash': 'stt',
    'qwen3-asr-flash-2025-09-08': 'stt',
    'qwen3-asr-flash-realtime': 'stt',
    'qwen3-asr-flash-realtime-2025-10-27': 'stt',
    'qwen2-audio-instruct': 'stt',
    'qwen3-livetranslate-flash-realtime': 'stt',
    'qwen3-livetranslate-flash-realtime-2025-09-22': 'stt',
    'gummy-realtime-v1': 'stt',
    'gummy-chat-v1': 'stt',
    'fun-asr': 'stt',
    'fun-asr-2025-08-25': 'stt',
    'fun-asr-mtl': 'stt',
    'fun-asr-mtl-2025-08-25': 'stt',
    'fun-asr-realtime': 'stt',
    'fun-asr-realtime-2025-09-15': 'stt',
    'paraformer-v2': 'stt',
    'paraformer-8k-v2': 'stt',
    'paraformer-v1': 'stt',
    'paraformer-8k-v1': 'stt',
    'paraformer-mtl-v1': 'stt',
    'paraformer-realtime-v2': 'stt',
    'paraformer-realtime-v1': 'stt',
    'paraformer-realtime-8k-v2': 'stt',
    'paraformer-realtime-8k-v1': 'stt',
    'sensevoice-v1': 'stt',

    # SD
    'qwen-image-plus': 'sd',
    'qwen-image': 'sd',
    'qwen-image-edit-plus': 'sd',
    'qwen-image-edit-plus-2025-10-30': 'sd',
    'qwen-image-edit': 'sd',
    'qwen-mt-image': 'sd',
    'wan2.5-t2i-preview': 'sd',
    'wan2.2-t2i-plus': 'sd',
    'wan2.2-t2i-flash': 'sd',
    'wanx2.1-t2i-plus': 'sd',
    'wanx2.1-t2i-turbo': 'sd',
    'wanx2.0-t2i-turbo': 'sd',
    'wanx-v1': 'sd',
    'wan2.5-i2i-preview': 'sd',
    'wanx2.1-imageedit': 'sd',
    'wanx-sketch-to-image-lite': 'sd',
    'wanx-x-painting': 'sd',
    'wanx-style-repaint-v1': 'sd',
    'wanx-background-generation-v2': 'sd',
    'image-out-painting': 'sd',
    'image-instance-segmentation': 'sd',
    'image-erase-completion': 'sd',
    'wanx-virtualmodel': 'sd',
    'virtualmodel-v2': 'sd',
    'shoemodel-v1': 'sd',
    'wanx-poster-generation-v1': 'sd',
    'facechain-facedetect': 'sd',
    'facechain-finetune': 'sd',
    'facechain-generation': 'sd',
    'wordart-texture': 'sd',
    'wordart-semantic': 'sd',
    'aitryon': 'sd',
    'aitryon-plus': 'sd',
    'aitryon-parsing-v1': 'sd',
    'aitryon-refiner': 'sd',
    'wan2.5-t2v-preview': 'sd',
    'wan2.2-t2v-plus': 'sd',
    'wanx2.1-t2v-turbo': 'sd',
    'wanx2.1-t2v-plus': 'sd',
    'wan2.5-i2v-preview': 'sd',
    'wan2.2-i2v-flash': 'sd',
    'wan2.2-i2v-plus': 'sd',
    'wanx2.1-i2v-turbo': 'sd',
    'wanx2.1-i2v-plus': 'sd',
    'wan2.2-kf2v-flash': 'sd',
    'wanx2.1-kf2v-plus': 'sd',
    'wanx2.1-vace-plus': 'sd',
    'wan2.2-s2v-detect': 'sd',
    'wan2.2-s2v': 'sd',
    'wan2.2-animate-move': 'sd',
    'wan2.2-animate-mix': 'sd',
    'animate-anyone-detect-gen2': 'sd',
    'animate-anyone-template-gen2': 'sd',
    'animate-anyone-gen2ss': 'sd',
    'animate-anyone-detect': 'sd',
    'animate-anyone': 'sd',
    'emo-detect-v1': 'sd',
    'emo-v1': 'sd',
    'emo-detect': 'sd',
    'emo': 'sd',
    'liveportrait-detect': 'sd',
    'liveportrait': 'sd',
    'emoji-detect-v1': 'sd',
    'emoji-v1': 'sd',
    'videoretalk': 'sd',

    # TTS
    'qwen3-tts-flash': 'tts',
    'qwen3-tts-flash-2025-09-18': 'tts',
    'qwen3-tts-flash-realtime': 'tts',
    'qwen3-tts-flash-realtime-2025-09-18': 'tts',
    'qwen-tts-realtime': 'tts',
    'qwen-tts-realtime-latest': 'tts',
    'qwen-tts-realtime-2025-07-15': 'tts',
    'cosyvoice-v3-plus': 'tts',
    'cosyvoice-v3': 'tts',
    'cosyvoice-v2': 'tts',
    'cosyvoice-v1': 'tts',

    # EMBED
    'text-embedding-v4': 'embed',
    'text-embedding-v3': 'embed',
    'text-embedding-v2': 'embed',
    'text-embedding-v1': 'embed',
    'text-embedding-async-v2': 'embed',
    'text-embedding-async-v1': 'embed',

    # CROSS_MODAL_EMBED
    'qwen2.5-vl-embedding': 'cross_modal_embed',
    'tongyi-embedding-vision-plus': 'cross_modal_embed',
    'tongyi-embedding-vision-flash': 'cross_modal_embed',
    'multimodal-embedding-v1': 'cross_modal_embed',

    # RERANK
    'qwen3-rerank': 'rerank',
    'gte-rerank-v2': 'rerank',

    # LLM (others)
    'tongyi-intent-detect-v3': 'llm',
    'qwen-plus-character': 'llm',

    # ===== Doubao =====
    'doubao-seed-1-6-lite-251015': 'vlm',
    'doubao-seed-1-6-vision-250815': 'vlm',
    'doubao-seed-1-6-250615': 'vlm',
    'doubao-seed-1-6-251015': 'vlm',
    'doubao-seed-1-6-flash-250828': 'vlm',
    'doubao-seed-1-6-thinking-250715': 'vlm',
    'doubao-seed-1-6-thinking-250615': 'vlm',
    'doubao-seed-1-6-flash-250715': 'vlm',
    'doubao-seed-1-6-flash-250615': 'vlm',
    'doubao-1-5-thinking-pro-250415': 'llm',
    'doubao-1-5-thinking-pro-m-250428': 'vlm',
    'doubao-1-5-vision-pro-250328': 'vlm',
    'doubao-1-5-pro-32k-250115': 'llm',
    'doubao-1-5-pro-32k-character-250228': 'llm',
    'doubao-1-5-pro-256k-250115': 'llm',
    'doubao-1-5-lite-32k-250115': 'llm',
    'doubao-1-5-ui-tars-250428': 'vlm',
    'doubao-1-5-thinking-vision-pro-250428': 'vlm',
    'doubao-1-5-vision-pro-32k-250115': 'vlm',
    'doubao-1-5-vision-lite-250315': 'vlm',
    'deepseek-v3-1-terminus': 'llm',
    'deepseek-v3-1-250821': 'llm',
    'doubao-seedance-1-0-pro-250528': 'sd',
    'doubao-seedance-1-0-pro-fast-251015': 'sd',
    'doubao-seedance-1-0-lite-t2v-250428': 'sd',
    'doubao-seedance-1-0-lite-i2v-250428': 'sd',
    'doubao-seedream-4-0-250828': 'sd',
    'doubao-seedream-3-0-t2i-250415': 'sd',
    'doubao-seededit-3-0-i2i-250628': 'sd',
    'doubao-seed3d-1-0-250928': 'sd',
    'doubao-embedding-large-text-250515': 'embed',
    'doubao-embedding-large-text-240915': 'embed',
    'doubao-embedding-text-240715': 'embed',
    'doubao-embedding-vision-250615': 'cross_modal_embed',
    'doubao-embedding-vision-250328': 'cross_modal_embed',

    # ===== DeepSeek =====
    'deepseek-chat': 'llm',
    'deepseek-reasoner': 'llm'
}
_TOKEN_MAP = {
    'embed': ('embedding', 'embed'),
    'stt': ('whisper', 'paraformer', 'asr', 'stt', 'transcribe'),
    'tts': ('tts', 'cosyvoice', 'nova-tts'),
    'vlm': ('qwen-vl', 'vl', 'vision', 'caption', 'omni', 'vlm', 'seed'),
    'ocr': ('ocr',),
    'rerank': ('rerank',),
    'cross_modal_embed': ('cross_modal', 'multimodal-embedding', 'embedding-vision'),
    'sd': ('dall', 'wan', 'sora', 'image', 'video', 't2i', 't2v'),
}
_SUFFIX_RE = re.compile(
    r'(?:'
    r'|[-_.]\d{4}(?:-\d{2}-\d{2})?'   # date-like suffix
    r'|[-_.]v?\d+[a-z0-9\-]*'         # version-like suffix
    r')+$',
    flags=re.I
)

def _normalize_key(name: str) -> str:
    '''Normalize model name for consistent comparison.'''
    if not name:
        return ''
    s = name.strip().lower()
    s = re.sub(r'[^a-z0-9]+', '-', s)
    s = s.strip('-')
    s = re.sub(_SUFFIX_RE, '', s)
    return s

def _contains_token(name: str, token: str) -> bool:
    '''Check if token exists as a separate unit in the model name.'''
    if not token:
        return False
    patterns = [
        rf'(^|[-_.]){re.escape(token)}($|[-_.])',
        rf'{re.escape(token)}$',
        rf'^{re.escape(token)}'
    ]
    return any(re.search(p, name) for p in patterns)

NORMALIZED_MODEL_MAPPING: Dict[str, str] = {
    _normalize_key(k): v for k, v in MODEL_MAPPING.items()
}

def feature_keyword_rule(model_name: str) -> Optional[str]:
    '''Identify the model type by normalized name or keyword match.'''
    stripped_input = _normalize_key(model_name)
    if stripped_input in NORMALIZED_MODEL_MAPPING:
        return NORMALIZED_MODEL_MAPPING[stripped_input]

    for model_type, tokens in _TOKEN_MAP.items():
        for tok in tokens:
            if _contains_token(stripped_input, tok):
                return model_type
    return None
def special_model_rule(model_name: str) -> Optional[str]:
    '''Determine the model category'''
    return MODEL_MAPPING.get(model_name)

def get_model_type(model_name: str) -> str:
    model_name = model_name.lower()
    if not model_name:
        return 'llm'
    for rule in (special_model_rule, feature_keyword_rule):
        try:
            result = rule(model_name)
            if result:
                LOG.info(f'Model: {model_name} classified as type: {result} by rule: {rule.__name__}')
                return result
        except Exception as e:
            LOG.warning(f'Rule {rule.__name__} failed: {e}')
    LOG.warning(f'Cannot classify model type for: {model_name}. Defaulting to "llm" instead.')
    return 'llm'
